import torch


def make_optimizer(cfg, model, center_criterion):
    """
    创建用于模型训练的优化器。

    Args:
        cfg (CfgNode): 包含训练配置的配置节点。
        model (torch.nn.Module): 待训练的模型。
        center_criterion (torch.nn.Module): 中心损失的模型。

    Returns:
        optimizer (torch.optim.Optimizer): 模型参数的优化器。
        optimizer_center (torch.optim.SGD): 中心损失参数的优化器。
    """
    params = []
    for key, value in model.named_parameters():
        if not value.requires_grad:
            continue
        lr = cfg.SOLVER.BASE_LR
        weight_decay = cfg.SOLVER.WEIGHT_DECAY
        if "bias" in key:
            lr = cfg.SOLVER.BASE_LR * cfg.SOLVER.BIAS_LR_FACTOR
            weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS
        if cfg.SOLVER.LARGE_FC_LR:
            if "classifier" in key or "arcface" in key:
                lr = cfg.SOLVER.BASE_LR * 2
                print('Using two times learning rate for fc ')
        params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]

    # 根据配置选择使用 SGD 还是 AdamW 优化器
    if cfg.SOLVER.OPTIMIZER_NAME == 'SGD':
        optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params, momentum=cfg.SOLVER.MOMENTUM)
    elif cfg.SOLVER.OPTIMIZER_NAME == 'Adam':
        optimizer = torch.optim.AdamW(params, lr=cfg.SOLVER.BASE_LR, weight_decay=cfg.SOLVER.WEIGHT_DECAY)
    else:
        optimizer = getattr(torch.optim, cfg.SOLVER.OPTIMIZER_NAME)(params)

    # 创建用于中心损失的 SGD 优化器
    optimizer_center = torch.optim.SGD(center_criterion.parameters(), lr=cfg.SOLVER.CENTER_LR)

    return optimizer, optimizer_center
